package org.jetlinks.pro.datasource.rdb;

import com.zaxxer.hikari.HikariDataSource;
import io.r2dbc.pool.ConnectionPool;
import io.r2dbc.pool.ConnectionPoolConfiguration;
import io.r2dbc.spi.Connection;
import io.r2dbc.spi.ConnectionFactory;
import org.apache.commons.collections4.MapUtils;
import org.hswebframework.ezorm.rdb.executor.reactive.ReactiveSyncSqlExecutor;
import org.hswebframework.ezorm.rdb.metadata.RDBDatabaseMetadata;
import org.hswebframework.ezorm.rdb.metadata.RDBSchemaMetadata;
import org.hswebframework.ezorm.rdb.operator.DatabaseOperator;
import org.hswebframework.ezorm.rdb.operator.DefaultDatabaseOperator;
import org.hswebframework.web.bean.FastBeanCopier;
import org.hswebframework.web.crud.sql.DefaultR2dbcExecutor;
import org.jetlinks.pro.datasource.AbstractDataSource;
import org.jetlinks.pro.datasource.Command;
import org.jetlinks.pro.datasource.DataSourceType;
import org.jetlinks.pro.datasource.rdb.command.RDBCommand;
import org.springframework.boot.autoconfigure.r2dbc.ConnectionFactoryBuilder;
import org.springframework.boot.autoconfigure.r2dbc.EmbeddedDatabaseConnection;
import org.springframework.boot.autoconfigure.r2dbc.R2dbcProperties;
import org.springframework.boot.context.properties.PropertyMapper;
import org.springframework.data.r2dbc.connectionfactory.ConnectionFactoryUtils;
import org.springframework.data.r2dbc.connectionfactory.R2dbcTransactionManager;
import org.springframework.transaction.ReactiveTransactionManager;
import org.springframework.transaction.reactive.TransactionalOperator;
import org.springframework.util.ClassUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.Closeable;
import java.io.IOException;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;

public class DefaultRDBDataSource extends AbstractDataSource<RDBDataSourceProperties> implements RDBDataSource {
    private DatabaseOperator operator;
    private final List<Closeable> closeables = new CopyOnWriteArrayList<>();

    private TransactionalOperator transactionalOperator;

    public DefaultRDBDataSource(String id,
                                RDBDataSourceProperties config) {
        super(id, config);
        init();
    }

    @Override
    public DatabaseOperator operator() {
        return operator;
    }

    @Override
    public DataSourceType getType() {
        return RDBDataSourceType.rdb;
    }

    public void init() {
        if (getConfig().getType() == RDBDataSourceProperties.Type.r2dbc) {
            initR2dbc();
        } else {
            initJdbc();
        }
    }

    public synchronized void initR2dbc() {
        R2dbcProperties properties = new R2dbcProperties();
        properties.setUrl(getConfig().getUrl());
        properties.setUsername(getConfig().getUsername());
        properties.setPassword(getConfig().getPassword());
        ConnectionFactory connectionFactory = ConnectionFactoryBuilder
            .of(properties, () -> EmbeddedDatabaseConnection.get(ClassUtils.getDefaultClassLoader()))
            .build();
        R2dbcProperties.Pool pool = properties.getPool();

        ConnectionPoolConfiguration.Builder builder = ConnectionPoolConfiguration.builder(connectionFactory);
        PropertyMapper map = PropertyMapper.get().alwaysApplyingWhenNonNull();
        map.from(pool.getMaxIdleTime()).to(builder::maxIdleTime);
        map.from(pool.getInitialSize()).to(builder::initialSize);
        map.from(pool.getMaxSize()).to(builder::maxSize);
        map.from(pool.getValidationQuery()).whenHasText().to(builder::validationQuery);

        builder.maxLifeTime(Duration.ofMinutes(10));
        builder.maxAcquireTime(Duration.ofSeconds(10));

        if (StringUtils.hasText(pool.getValidationQuery())) {
            builder.validationQuery(pool.getValidationQuery());
        }
        ConnectionPool connectionPool = new ConnectionPool(builder.build());
        closeables.add(connectionPool);
        transactionalOperator = TransactionalOperator.create(new R2dbcTransactionManager(connectionPool));

        DefaultR2dbcExecutor executor = new DefaultR2dbcExecutor() {
            @Override
            protected Mono<Connection> getConnection() {
                return ConnectionFactoryUtils.getConnection(connectionPool);
            }
        };
        executor.setBindSymbol(getConfig().getDialect().getBindSymbol());
        executor.setBindCustomSymbol(!executor.getBindSymbol().equals("?"));
        RDBDatabaseMetadata database = new RDBDatabaseMetadata(getConfig().getDialect().getDialect());
        database.addFeature(executor);
        database.addFeature(ReactiveSyncSqlExecutor.of(executor));

        RDBSchemaMetadata schema = getConfig().getDialect().createSchema(getConfig().getSchema());
        database.addSchema(schema);
        database.setCurrentSchema(schema);
        this.operator = DefaultDatabaseOperator.of(database);

    }

    public synchronized void initJdbc() {
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setJdbcUrl(getConfig().getUrl());
        dataSource.setUsername(getConfig().getUsername());
        dataSource.setPassword(getConfig().getPassword());
        if (MapUtils.isNotEmpty(getConfig().getOthers())) {
            FastBeanCopier.copy(getConfig().getOthers(), dataSource);
        }
        closeables.add(dataSource);
        RDBDatabaseMetadata database = new RDBDatabaseMetadata(getConfig().getDialect().getDialect());
        database.addFeature(new RDBJdbcReactiveSqlExecutor(dataSource));
        database.addFeature(new RDBJdbcSyncSqlExecutor(dataSource));

        RDBSchemaMetadata schema = getConfig().getDialect().createSchema(getConfig().getSchema());
        database.addSchema(schema);
        database.setCurrentSchema(schema);

        this.operator = DefaultDatabaseOperator.of(database);
        this.transactionalOperator = null;
    }

    @Override
    @SuppressWarnings("all")
    protected <R> R handleNoHandlerCommand(Command<R> command) {
        if (command instanceof RDBCommand) {
            R r = ((RDBCommand<R>) command).execute(operator);
            if (transactionalOperator != null) {
                if (r instanceof Mono) {
                    return (R) transactionalOperator.transactional(((Mono) r));
                } else if (r instanceof Flux) {
                    return (R) transactionalOperator.transactional(((Flux) r));
                }
            }
            return r;
        }
        return super.handleNoHandlerCommand(command);
    }

    @Override
    protected void handleSetConfig(RDBDataSourceProperties oldConfig,
                                   RDBDataSourceProperties newConfig) {
        if (oldConfig == null) {
            return;
        }
        releaseOld();
        init();
    }

    protected void releaseOld() {
        closeables.removeIf(closeable -> {
            try {
                closeable.close();
            } catch (IOException ignore) {
            }
            return true;
        });
    }

    @Override
    protected void doOnDispose() {
        releaseOld();
    }
}
